import sys

import pygame
import numpy as np


class Chess():
    def __init__(self):

        self.GRID_SIZE = 8
        self.CELL_SIZE = 30
        self.MARGIN = self.CELL_SIZE
        self.WINDOW_SIZE = (self.GRID_SIZE - 1) * self.CELL_SIZE + 2 * self.MARGIN

        # 定义颜色
        self.WHITE = (255, 255, 255)
        self.GREEN = (187, 255, 170)
        self.BLACK = (0, 0, 0)
        self.RED = (200, 0, 0)
        self.TRANSPARENT = (0, 0, 0, 0)

        pygame.init()
        self.screen = pygame.display.set_mode((self.WINDOW_SIZE, self.WINDOW_SIZE))
        pygame.display.set_caption("五指棋大战AI")

        self.board = np.zeros((self.GRID_SIZE, self.GRID_SIZE))

        # import
        self.init_board()

    def init_board(self):
        self.ai_board = Board(width=self.GRID_SIZE, height=self.GRID_SIZE)
        self.ai_board.init_board(1)
        self.board = np.zeros((self.GRID_SIZE, self.GRID_SIZE))

        self.current_pos = None
        self.win_group = None

    def draw_board(self):
        # 绘制最大的黑色框
        pygame.draw.rect(self.screen, self.BLACK, (
            self.MARGIN - 1, self.MARGIN - 1, (self.GRID_SIZE - 1) * self.CELL_SIZE + 2,
            (self.GRID_SIZE - 1) * self.CELL_SIZE + 2),
                         1)

        if self.current_pos is not None:
            # 计算矩形的位置
            current_x = self.MARGIN + self.current_pos[1] * self.CELL_SIZE
            current_y = self.MARGIN + self.current_pos[0] * self.CELL_SIZE
            pygame.draw.circle(self.screen, self.RED, (current_x, current_y), self.CELL_SIZE // 2.5)

        if self.win_group is not None:
            for item in self.win_group:
                # 计算矩形的位置
                current_x = self.MARGIN + item[1] * self.CELL_SIZE
                current_y = self.MARGIN + item[0] * self.CELL_SIZE
                pygame.draw.circle(self.screen, self.RED, (current_x, current_y), self.CELL_SIZE // 2)

        for row in range(self.GRID_SIZE):
            for col in range(self.GRID_SIZE):
                # 计算矩形的位置
                rect_x = self.MARGIN + col * self.CELL_SIZE
                rect_y = self.MARGIN + row * self.CELL_SIZE

                if row != self.GRID_SIZE - 1 and col != self.GRID_SIZE - 1:
                    # 绘制黑色边框
                    pygame.draw.rect(self.screen, self.BLACK, (rect_x, rect_y, self.CELL_SIZE, self.CELL_SIZE), 1)

                # 如果该单元格有棋子，则绘制
                if self.board[row][col] != 0:
                    item = self.board[row][col]
                    pygame.draw.circle(self.screen, self.BLACK if item == 1 else self.WHITE,
                                       (rect_x, rect_y),
                                       self.CELL_SIZE // 3)

    def down_piece(self, pos, user) -> bool:
        # 获取鼠标点击的位置
        row, col = pos
        if self.board[row, col] == 0:
            self.board[row, col] = user
            self.current_pos = pos
            a = self.location_to_move(row, col)
            self.ai_board.do_move(a)
            return True
        else:
            return False

    def get_board(self):
        return self.board

    def fill(self):
        self.screen.fill(self.GREEN)

    def set_win(self, group):
        self.win_group = group

    def get_ai_board(self):
        return self.ai_board

    def location_to_move(self, *location):
        if len(location) != 2:
            return -1
        h = location[0]
        w = location[1]
        move = h * self.GRID_SIZE + w
        if move not in range(self.GRID_SIZE * self.GRID_SIZE):
            return -1
        return move


class Board(object):
    """board for the game"""

    def __init__(self, **kwargs):
        self.width = int(kwargs.get('width', 8))
        self.height = int(kwargs.get('height', 8))
        # board states stored as a dict,
        # key: move as location on the board,
        # value: player as pieces type
        self.states = {}
        # need how many pieces in a row to win
        self.n_in_row = int(kwargs.get('n_in_row', 5))
        self.players = [1, 2]  # player1 and player2

    def init_board(self, start_player=0):
        if self.width < self.n_in_row or self.height < self.n_in_row:
            raise Exception('board width and height can not be '
                            'less than {}'.format(self.n_in_row))
        self.current_player = self.players[start_player]  # start player
        # keep available moves in a list
        self.availables = list(range(self.width * self.height))
        self.states = {}
        self.last_move = -1

    def move_to_location(self, move):
        """
        3*3 board's moves like:
        6 7 8
        3 4 5
        0 1 2
        and move 5's location is (1,2)
        """
        h = move // self.width
        w = move % self.width
        return [h, w]

    def location_to_move(self, location):
        if len(location) != 2:
            return -1
        h = location[0]
        w = location[1]
        move = h * self.width + w
        if move not in range(self.width * self.height):
            return -1
        return move

    def current_state(self):
        """return the board state from the perspective of the current player.
        state shape: 4*width*height
        """

        square_state = np.zeros((4, self.width, self.height))
        if self.states:
            moves, players = np.array(list(zip(*self.states.items())))
            move_curr = moves[players == self.current_player]
            move_oppo = moves[players != self.current_player]
            square_state[0][move_curr // self.width,
                            move_curr % self.height] = 1.0
            square_state[1][move_oppo // self.width,
                            move_oppo % self.height] = 1.0
            # indicate the last move location
            square_state[2][self.last_move // self.width,
                            self.last_move % self.height] = 1.0
        if len(self.states) % 2 == 0:
            square_state[3][:, :] = 1.0  # indicate the colour to play
        return square_state[:, ::-1, :]

    def do_move(self, move):
        self.states[move] = self.current_player
        self.availables.remove(move)
        self.current_player = (
            self.players[0] if self.current_player == self.players[1]
            else self.players[1]
        )
        self.last_move = move

    def has_a_winner(self):
        width = self.width
        height = self.height
        states = self.states
        n = self.n_in_row

        moved = list(set(range(width * height)) - set(self.availables))
        if len(moved) < self.n_in_row * 2 - 1:
            return False, -1

        for m in moved:
            h = m // width
            w = m % width
            player = states[m]

            if (w in range(width - n + 1) and
                    len(set(states.get(i, -1) for i in range(m, m + n))) == 1):
                return True, player

            if (h in range(height - n + 1) and
                    len(set(states.get(i, -1) for i in range(m, m + n * width, width))) == 1):
                return True, player

            if (w in range(width - n + 1) and h in range(height - n + 1) and
                    len(set(states.get(i, -1) for i in range(m, m + n * (width + 1), width + 1))) == 1):
                return True, player

            if (w in range(n - 1, width) and h in range(height - n + 1) and
                    len(set(states.get(i, -1) for i in range(m, m + n * (width - 1), width - 1))) == 1):
                return True, player

        return False, -1

    def game_end(self):
        """Check whether the game is ended or not"""
        win, winner = self.has_a_winner()
        if win:
            return True, winner
        elif not len(self.availables):
            return True, -1
        return False, -1

    def get_current_player(self):
        return self.current_player
